package com.kingsoft.dc.khaos.extender.auth.utils;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLUseStatement;
import com.alibaba.druid.sql.parser.ParserException;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.druid.stat.TableStat.Name;
import com.kingsoft.dc.khaos.extender.meta.model.AuthTable;
import com.kingsoft.dc.khaos.extender.model.DDLType;
import com.kingsoft.dc.khaos.extender.model.DMLType;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.stream.Collectors;

/**
 * sql解析工具类.
 */
public class SqlParser {

    private static Logger logger = LoggerFactory.getLogger(SqlParser.class);
    /**
     * 提取sql要操作的表，及操作类型
     *
     * @param sql 待分析的sql，支持多条sql
     * @return
     * @throws ParserException
     */
    public static Map<String, TreeSet<String>> getAuthTables(String sql, String dbType) throws ParserException {
        String dbTypeLower = dbType.toLowerCase();
        System.out.println("数据类型是===> :" + dbTypeLower);
        if (dbTypeLower.equalsIgnoreCase("mpp") || dbTypeLower.equalsIgnoreCase("hawq") || dbTypeLower.equalsIgnoreCase("greenplum") || dbTypeLower.equalsIgnoreCase("hashdata")) {
            dbTypeLower = "postgresql";
        }
        System.out.println("数据类型是===> :" + dbTypeLower);
        System.out.println("执行sql是===> :" + sql);
        List<SQLStatement> stmts = SQLUtils.parseStatements(sql, dbTypeLower);
        if (stmts == null) {
            return null;
        }

        TreeSet<AuthTable> authTables = new TreeSet<>();
        Map<String, TreeSet<String>> tableAuthMap = new HashMap<>();
        HashSet<String> contextTableSet = new HashSet<>(); //存放临时表的上下文集合
        String database = "DEFAULT";
        for (SQLStatement stmt : stmts) {
            SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(dbTypeLower);
            if (stmt instanceof SQLUseStatement) {
                database = ((SQLUseStatement) stmt).getDatabase().getSimpleName();
                continue;
            }
            stmt.accept(statVisitor);
            Map<Name, TableStat> tables = statVisitor.getTables();
            if (tables != null) {

                final String db = database;

                tables.forEach((tableName, stat) -> {
                    /*if (stat.getTempTableCount() > 0 ) { // 跳过mpp创建临时表权限
                        return;
                     }*/
                    String tableFullName = tableName.getName();
                    System.out.println("tableFullName:" + tableFullName);
                    //企业云2.2配合数据管理修改oracle和mpp大小写敏感问题
                    if (!tableFullName.contains(".")) {
                        tableFullName = db + "." + tableFullName;
                        System.out.println("tableFullName！.:" + tableFullName);
                    } else if (!tableFullName.contains(".") && (dbType.toLowerCase().equalsIgnoreCase("oracle") || dbType.toLowerCase().equalsIgnoreCase("mpp"))) {
                        tableFullName = db + "." + tableFullName.substring(1, tableFullName.length() - 1);
                        System.out.println("tableFullName==========>:    " + tableFullName);
                    } else if (tableFullName.contains(".") && (dbType.toLowerCase().equalsIgnoreCase("oracle") || dbType.toLowerCase().equalsIgnoreCase("mpp") || dbType.toLowerCase().equalsIgnoreCase("greenplum") || dbType.toLowerCase().equalsIgnoreCase("hawq") || dbType.toLowerCase().equalsIgnoreCase("hashdata"))) {
                        String[] dbAndTable = tableFullName.split("\\.");
                        if (dbAndTable[0].charAt(0) == '"' && dbAndTable[0].charAt(dbAndTable[0].length() - 1) == '"' && dbAndTable[1].charAt(0) == '"' && dbAndTable[1].charAt(dbAndTable[1].length() - 1) == '"') {
                            tableFullName = dbAndTable[0].substring(1, dbAndTable[0].length() - 1) + "." + dbAndTable[1].substring(1, dbAndTable[1].length() - 1);
                            System.out.println("tableFullName==========>:    " + tableFullName);
                        }
                    }
                    // 把临时表加入上下文中
                    if (stat.getTempTableCount() > 0) {
                        contextTableSet.add(tableFullName);
                    }
                    if (!tableAuthMap.containsKey(tableFullName)) {
                        tableAuthMap.put(tableFullName, new TreeSet<String>());
                    }
                    String priv = "";
                    if (stat.getInsertCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.INSERT.name());
                    }
                    if (stat.getSelectCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.SELECT.name());
                    }
                    if (stat.getCreateCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.CREATE.name());
                    }
                    if (stat.getDeleteCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.DELETE.name());
                    }
                    if (stat.getUpdateCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.UPDATE.name());
                    }
                    if (stat.getDropCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.DROP.name());
                    }
                    if (stat.getAlterCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.ALTER.name());
                    }
                    if (stat.getCreateIndexCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.CREATE_INDEX.name());
                    }
                    if (stat.getDropIndexCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.DROP_INDEX.name());
                    }
                });
            }
        }

        logger.info(String.format("filterBeforeTableAuthMap:[%s]", tableAuthMap.keySet().stream().collect(Collectors.joining(","))));
        // 在鉴权集合tableAuthMap中删除掉临时表及action
        for (String tmpTableName : contextTableSet) {
            tableAuthMap.keySet().removeIf(key -> key.equalsIgnoreCase(tmpTableName));
        }
        logger.info(String.format("filterAfterTableAuthMap:[%s]", tableAuthMap.keySet().stream().collect(Collectors.joining(","))));
        return tableAuthMap;
    }

    //mppsql插件和sparksql插件，增加dbName
    public static Map<String, TreeSet<String>> getAuthTables(String sql, String dbType, String dbName,List<AuthTable> viewAuthTables) throws ParserException {
        String dbTypeLower = dbType.toLowerCase();
        System.out.println("spark数据类型 :" + dbTypeLower);    //hive
        if (dbTypeLower.equalsIgnoreCase("mpp")) {
            dbTypeLower = "postgresql";
        }
        System.out.println("spark数据类型是 :" + dbTypeLower);    //hive
        System.out.println("spark执行sql是 :" + sql);   //select * from  where id = 2
        List<SQLStatement> stmts = SQLUtils.parseStatements(sql, dbTypeLower);
        if (stmts == null) {
            return null;
        }

        TreeSet<AuthTable> authTables = new TreeSet<>();
        Map<String, TreeSet<String>> tableAuthMap = new HashMap<>();
        HashSet<String> contextTableSet = new HashSet<>(); //存放临时表的上下文集合
        String database = "DEFAULT";
        for (SQLStatement stmt : stmts) {
            SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(dbTypeLower);
            if (stmt instanceof SQLUseStatement) {
                database = ((SQLUseStatement) stmt).getDatabase().getSimpleName();
                continue;
            }
            stmt.accept(statVisitor);
            Map<Name, TableStat> tables = statVisitor.getTables();
            String db = database;
            if (tables != null) {
                for (Map.Entry<Name,TableStat> entry : tables.entrySet()) {
                    if (StringUtils.isNotBlank(dbName)) {
                        System.out.println("dbName不为空");
                        db = dbName;
                    }
                    Name tableName = entry.getKey();
                    TableStat stat = entry.getValue();
                    if (null != viewAuthTables && viewAuthTables.size() > 0) {
                        System.out.println("viewAuthTables :" + viewAuthTables.toString());
                        for (AuthTable viewAuthTable : viewAuthTables) {
                            System.out.println("SqlTable:" + tableName.getName());
                            System.out.println("AuthTable:" + viewAuthTable.getTblName());
                            if (tableName.getName().equals(viewAuthTable.getTblName())) {
                                db = viewAuthTable.getDbName();
                            }
                        }
                    }
                    String tableFullName = tableName.getName();
                    System.out.println("down_tableFullNameDbName:" + tableFullName);//temp_view视图表
                    if (!tableFullName.contains(".")) {
                        tableFullName = db + "." + tableFullName;
                        System.out.println("down_tableFullName！.:" + tableFullName);
                    } else if (!tableFullName.contains(".") && (dbType.toLowerCase().equalsIgnoreCase("oracle") || dbType.toLowerCase().equalsIgnoreCase("mpp"))) {
                        tableFullName = db + "." + tableFullName.substring(1, tableFullName.length() - 1);
                        System.out.println("tableFullName==========>:    " + tableFullName);
                    } else if (tableFullName.contains(".") && (dbType.toLowerCase().equalsIgnoreCase("oracle") || dbType.toLowerCase().equalsIgnoreCase("mpp") || dbType.toLowerCase().equalsIgnoreCase("greenplum") || dbType.toLowerCase().equalsIgnoreCase("hawq") || dbType.toLowerCase().equalsIgnoreCase("hashdata"))) {
                        String[] dbAndTable = tableFullName.split("\\.");
                        if (dbAndTable[0].charAt(0) == '"' && dbAndTable[0].charAt(dbAndTable[0].length() - 1) == '"' && dbAndTable[1].charAt(0) == '"' && dbAndTable[1].charAt(dbAndTable[1].length() - 1) == '"') {
                            tableFullName = dbAndTable[0].substring(1, dbAndTable[0].length() - 1) + "." + dbAndTable[1].substring(1, dbAndTable[1].length() - 1);
                            System.out.println("tableFullName==========>:    " + tableFullName);
                        }
                    }
                    // 把临时表加入上下文中
                    if (stat.getTempTableCount() > 0) {
                        contextTableSet.add(tableFullName);
                    }
                    if (!tableAuthMap.containsKey(tableFullName)) {
                        tableAuthMap.put(tableFullName, new TreeSet<String>());
                    }
                    String priv = "";
                    if (stat.getInsertCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.INSERT.name());
                    }
                    if (stat.getSelectCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.SELECT.name());
                    }
                    if (stat.getCreateCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.CREATE.name());
                    }
                    if (stat.getDeleteCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.DELETE.name());
                    }
                    if (stat.getUpdateCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DMLType.UPDATE.name());
                    }
                    if (stat.getDropCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.DROP.name());
                    }
                    if (stat.getAlterCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.ALTER.name());
                    }
                    if (stat.getCreateIndexCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.CREATE_INDEX.name());
                    }
                    if (stat.getDropIndexCount() > 0) {
                        tableAuthMap.get(tableFullName).add(DDLType.DROP_INDEX.name());
                    }
                }
            }
        }

        logger.info(String.format("filterBeforeTableAuthMap:[%s]", tableAuthMap.keySet().stream().collect(Collectors.joining(","))));
        // 在鉴权集合tableAuthMap中删除掉临时表及action
        for (String tmpTableName : contextTableSet) {
            tableAuthMap.keySet().removeIf(key -> key.equalsIgnoreCase(tmpTableName));
        }
        logger.info(String.format("filterAfterTableAuthMap:[%s]", tableAuthMap.keySet().stream().collect(Collectors.joining(","))));
        return tableAuthMap;
    }

    public static void main(String[] args) {
//        String sql = "insert into tbl1 (id ,name) select id,name from sel_tab1; into insert_tbl2(id,name) values (1,';');insert into insert_tbl3(id,name) values (1,',')";
        String sql0 = "use aa;";
        String sql1 = "insert into table db1.tablename1 select a,b,c FROM (select a, b, c from db2.tablename2) ali;";
        String sql2 = "insert into table db11.tablename1 VALUES (1, ';');";
        String sql3 = "update table1 set id = 1 where id in (select id from table2) and name in (select id from table3);";
        String sql4 = "update table1 set id = 1";
        List<String> sb = new ArrayList<>();
        sb.add(sql0);
        sb.add(sql1);
        sb.add(sql2);
        sb.add(sql3);
        sb.add(sql4);
        String sql = StringUtils.join(sb.toArray(), ";");
    }

}
